# -*- coding: utf-8 -*-
"""
    walle-web

    :copyright: © 2015-2020 walle-web.io
    :created time: 2020-11-03 13:57:00
    :author: wen2013@163.com
"""
from threading import Thread
from flask import current_app
from flask_login import current_user
from flask_socketio import emit, Namespace
from walle.service.webssh.ssh import SSH
from walle.model.server import ServerModel


class WebSSHSocketIO(Namespace):
    namespace, room, app = None, None, None
    host, port, user, pkey, chan, ssh = None, None, None, None, None, None
    server_info = None
    ssh_start_flag = False

    def __init__(self, namespace, room=None, app=None):
        super(WebSSHSocketIO, self).__init__(namespace=namespace)
        self.room = room
        self.app = app

    def init_app(self, app):
        self.app = app

    def on_open(self, message):
        current_app.logger.info(message)
        if not current_user.is_authenticated:
            emit('close', {'event': 'disconnect', 'data': '未授权！'}, room=self.room)
            return

        if 'id' not in message:
            emit('close', {'event': 'disconnect', 'data': '请求信息不正确！'}, room=self.room)
            return
        server_model = ServerModel(id=message['id'])
        server_info = server_model.item()
        if not server_info:
            emit('close', {'event': 'disconnect', 'data': '找不到主机信息！'}, room=self.room)
            return
        if not server_info['enable_term']:
            emit('close', {'event': 'disconnect', 'data': '没有终端命令执行权限！'}, room=self.room)
            return
        self.server_info = server_info
        self.host = server_info['host']
        self.user = server_info['user']
        self.port = server_info['port']
        if self.ssh_start_flag:
            self.close_ssh()
        self.init_ssh()

    def on_term(self, message):
        current_app.logger.info(message)
        self.chan.send(message)

    def on_resize(self, message):
        current_app.logger.info(message)
        resize = message
        if resize and len(resize) == 2:
            self.chan.resize_pty(*resize)

    def on_disconnect(self):
        current_app.logger.info('{} Connection close'.format(self.host))
        self.close_ssh()

    def init_ssh(self):
        emit('server-info', {'event': 'server-info', 'data': self.server_info})
        emit('term', 'Connecting console terminal...\r\n')
        if not self.host:
            emit('term', 'Unknown host\r\n')
            self.close()
        try:
            self.ssh_start_flag = True
            self.ssh = SSH(self.host, self.port, self.user, self.pkey).get_client()
            self.chan = self.ssh.invoke_shell(term='xterm')
            self.chan.transport.set_keepalive(30)
            Thread(target=self.loop_read_ssh, args=(self.app,)).start()
        except Exception as e:
            current_app.logger.exception(e)
            emit('term', 'Exception: {}\r\n'.format(e))
            self.ssh_start_flag = False
            self.close_ssh()
            self.close()
            return

    def loop_read_ssh(self, app):
        self.emit('term-connected', 'Console terminal connected!\r\n')
        with app.app_context():
            while self.ssh_start_flag:
                data = self.chan.recv(32 * 1024)
                if not data:
                    self.emit('close', {'event': 'console-close'})
                    self.close_ssh()
                    break
                self.emit('term', data)

    def close_ssh(self):
        current_app.logger.info('{} ssh term close'.format(self.host))
        self.ssh_start_flag = False
        try:
            if self.chan:
                self.chan.close()
            if self.ssh:
                self.ssh.close()
        except Exception as e:
            current_app.logger.exception(e)
            self.emit('term', 'Exception: {}\r\n'.format(e))
            return


